import cv2
import numpy as np
import random
from rknnlite.api import RKNNLite
from enum import Enum
from PIL import Image, ImageDraw, ImageFont
from pathlib import Path

from utils import linear
from .utils.coco_utils import COCO_test_helper
from .classes import CLASSES


# 下面的是固定参数
current_dir = Path(__file__).parent

# 下面是可配置参数
OBJ_THRESH = 0.6
NMS_THRESH = 0.45

IMG_SIZE = (640, 640)  # (width, height), such as (1280, 736)

# 后面就不用动了

with open(current_dir / 'anchors.txt', 'r') as f:
    ANCHORS = np.array([float(_v) for _v in f.readlines()]).reshape(3,-1,2).tolist()

co_helper = COCO_test_helper(enable_letter_box=True)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def filter_boxes(boxes, box_confidences, box_class_probs):
    """Filter boxes with object threshold.
    """
    box_confidences = box_confidences.reshape(-1)
    candidate, class_num = box_class_probs.shape

    class_max_score = np.max(box_class_probs, axis=-1)
    classes = np.argmax(box_class_probs, axis=-1)

    if class_num==1:
        _class_pos = np.where(box_confidences >= OBJ_THRESH)
        scores = (box_confidences)[_class_pos]
    else:
        _class_pos = np.where(class_max_score* box_confidences >= OBJ_THRESH)
        scores = (class_max_score* box_confidences)[_class_pos]

    boxes = boxes[_class_pos]
    classes = classes[_class_pos]

    return boxes, classes, scores

def nms_boxes(boxes, scores):
    """Suppress non-maximal boxes.
    # Returns
        keep: ndarray, index of effective boxes.
    """
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    areas = w * h
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep

def box_process(position, anchors):
    grid_h, grid_w = position.shape[2:4]
    col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
    col = col.reshape(1, 1, grid_h, grid_w)
    row = row.reshape(1, 1, grid_h, grid_w)
    grid = np.concatenate((col, row), axis=1)
    stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1)

    # output format: xywh -> xyxy
    col = col.repeat(len(anchors), axis=0)
    row = row.repeat(len(anchors), axis=0)
    anchors = np.array(anchors)
    anchors = anchors.reshape(*anchors.shape, 1, 1)

    box_xy = position[:,:2,:,:]*2 - 0.5
    box_wh = pow(position[:,2:4,:,:]*2, 2) * anchors

    box_xy += grid
    box_xy *= stride
    box = np.concatenate((box_xy, box_wh), axis=1)

    # Convert [c_x, c_y, w, h] to [x1, y1, x2, y2]
    xyxy = np.copy(box)
    xyxy[:, 0, :, :] = box[:, 0, :, :] - box[:, 2, :, :]/ 2  # top left x
    xyxy[:, 1, :, :] = box[:, 1, :, :] - box[:, 3, :, :]/ 2  # top left y
    xyxy[:, 2, :, :] = box[:, 0, :, :] + box[:, 2, :, :]/ 2  # bottom right x
    xyxy[:, 3, :, :] = box[:, 1, :, :] + box[:, 3, :, :]/ 2  # bottom right y

    return xyxy

def post_process(input_data, anchors):
    boxes, scores, classes_conf = [], [], []
    # 1*255*h*w -> 3*85*h*w
    input_data = [_in.reshape([len(anchors[0]),-1]+list(_in.shape[-2:])) for _in in input_data]
    for i in range(len(input_data)):
        boxes.append(box_process(input_data[i][:,:4,:,:], anchors[i]))
        scores.append(input_data[i][:,4:5,:,:])
        classes_conf.append(input_data[i][:,5:,:,:])

    def sp_flatten(_in):
        ch = _in.shape[1]
        _in = _in.transpose(0,2,3,1)
        return _in.reshape(-1, ch)

    boxes = [sp_flatten(_v) for _v in boxes]
    classes_conf = [sp_flatten(_v) for _v in classes_conf]
    scores = [sp_flatten(_v) for _v in scores]

    boxes = np.concatenate(boxes)
    classes_conf = np.concatenate(classes_conf)
    scores = np.concatenate(scores)

    # filter according to threshold
    boxes, classes, scores = filter_boxes(boxes, scores, classes_conf)

    # nms
    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]
        keep = nms_boxes(b, s)

        if len(keep) != 0:
            nboxes.append(b[keep])
            nclasses.append(c[keep])
            nscores.append(s[keep])

    if not nclasses and not nscores:
        return None, None, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)

    return boxes, classes, scores

class Detecter:

    class MODEL_NPU_CORE(Enum):
        CORE_AUTO = RKNNLite.NPU_CORE_AUTO
        CORE_0 = RKNNLite.NPU_CORE_0
        CORE_1 = RKNNLite.NPU_CORE_1
        CORE_2 = RKNNLite.NPU_CORE_2
    
    @staticmethod
    def init(config):
        config.obj_thresh = linear(0.25, 0.8, 1 - config.alert_sensitivity)

    # model如果为None，则表示不加载模型，纯用绘图功能，不能用detect
    def __init__(self, config, model_npu: MODEL_NPU_CORE=None):
        global OBJ_THRESH
        OBJ_THRESH = config.obj_thresh

        self._model = RKNNLite()
        self._model.load_rknn(current_dir / 'model.rknn')
        if model_npu is not None:
            self._model.init_runtime(core_mask=model_npu.value)

    # 不会改变原图
    def detect(self, image):
        img, ratio, (dw, dh) = co_helper.letter_box(im=image, new_shape=(IMG_SIZE[1], IMG_SIZE[0]), pad_color=(0,0,0), info_need=True)

        outputs = self._model.inference([cv2.cvtColor(img, cv2.COLOR_BGR2RGB)])
        boxes, classes, scores = post_process(outputs, ANCHORS)

        if boxes is not None:
            for i in range(len(boxes)):
                bbox = boxes[i]
                bbox[0] -= dw
                bbox[1] -= dh
                bbox[2] -= dw
                bbox[3] -= dh
                boxes[i] = [value/ratio for value in bbox]
            
        return [(
            tuple(map(int, item[0])),  # box
            item[1],  # name
            item[2]  # score
        ) for item in zip(boxes, classes, scores)] if boxes is not None else []
